import socket, requests,os,tempfile
import urllib.request
from urllib.parse import urlparse
from ftplib import FTP
from enum import Enum
from support.utils.toolkit import Log

class NetHelper:
    CHUNK_SIZE = 8192
    ANONYMOUS = 'anonymous'
    PORT_MAP = Enum('PORT', {
    'http'  : 80,
    'https' : 443,
    'ftp'   : 21
    })
    
    STATUS_CODE_MAP = Enum('CODE', {
    'SUCCESS' : 200
    })
    
    def __init__(self):
        self.tmpdir = tempfile.TemporaryDirectory('analysis', 'nethelper')
            
    def __checkProtocol__(self, url):
        ret = True
        s = socket.socket()
        try:
            s.connect((url.netloc, NetHelper.PORT_MAP[url.scheme].value))
        except Exception as e:
            Log.e("Exception happend when we check protocol")
            ret = False
        finally:
            s.close()
        return ret
    
    def __downloadByHttp__(self, url, localdir, filename):
        ret = False
        response = requests.get(url, stream=True)
        if response.status_code == NetHelper.STATUS_CODE_MAP['SUCCESS'].value: 
            with open(os.path.join(localdir, filename), 'wb') as f: 
                for chunk in response.iter_content(chunk_size=NetHelper.CHUNK_SIZE):
                    f.write(chunk)
            ret = True
        else:
            Log.e(f"Web Server has err {response.status_code}")
        return ret
        
    def __downloadByFtp__(host, remotePath, localdir, filename, *, username=None, passwd=None):
        ret = False
        with FTP(host) as ftp:
            if(username is None):
                username = NetHelper.ANONYMOUS
            if(passwd is None):
                passwd = NetHelper.ANONYMOUS
            ftp.login(username, passwd)
            ftp.encoding = 'utf-8'
            with open(os.path.join(localdir, filename), 'wb') as f:
                ftp.retrbinary(f'RETR {remotePath}', f.write)
            ret = True
        return ret
    
    def __parseFtpUrl__(self, urlresult):
        username = None
        passwd = None
        host = None
        emt = urlresult.netloc.split('@')
        if(len(emt) > 1):
            info = emt[0].split(':')
            username = account[0]
            passwd = account[1]
            host = urlresult.scheme + "://" + emt[1]
        else:
            host = urlresult.scheme + "://" + emt[0]
        return username, passwd, host
        

    def download(self, url, *, targetDir=None):
        result = urlparse(url)
        finished = False
        if(targetDir is None):
            targetDir = self.tmpdir.name
        Log.i(f"we will download from {url} to {targetDir}")
        if(self.__checkProtocol__(result)):
            emt = result.path.split('/')
            filename = emt[-1]
            try:
                if(result.scheme in ['https', 'http']):
                    finished = self.__downloadByHttp__(url, targetDir, filename)
                if(result.scheme in ['ftp', 'ftps']):
                    username, passwd, host = self.__parseFtpUrl(result)
                    finished = self.__downloadByFtp__(host, result.path, targetDir, filename, username=username, passwd=passwd)
            except Exception as e:
                Log.e("Download has Err")
        if(finished):
            return os.path.join(targetDir, filename)
        return None
        
